# import sys
from prettytable import PrettyTable
# from sko.AFSA import AFSA
from streamlit_extras.colored_header import colored_header

from business.algorithm.utils import *


def run():
    colored_header(label="机器学习：模型推理", description=" ", color_name="violet-90")
    file = st.file_uploader("Upload `.csv`file", label_visibility="collapsed", accept_multiple_files=True)
    if len(file) < 2:
        table = PrettyTable(['file name', 'class', 'description'])
        table.add_row(['file_1', 'data set (+test data)', 'data file'])
        table.add_row(['file_2', 'model', 'model'])
        st.write(table)
    elif len(file) == 2:
        df = pd.read_csv(file[0])
        model_file = file[1]

        check_string_NaN(df)

        colored_header(label="数据信息", description=" ", color_name="violet-70")
        nrow = st.slider("rows", 1, len(df), 5)
        df_nrow = df.head(nrow)
        st.write(df_nrow)

        colored_header(label="特征&目标", description=" ", color_name="violet-70")

        target_num = st.number_input('目标数量', min_value=1, max_value=10, value=1)

        col_feature, col_target = st.columns(2)
        # features
        features = df.iloc[:, :-target_num]
        # targets
        targets = df.iloc[:, -target_num:]
        with col_feature:
            st.write(features.head())
        with col_target:
            st.write(targets.head())
        colored_header(label="target", description=" ", color_name="violet-70")

        target_selected_option = st.selectbox('target', list(targets)[::-1])

        targets = targets[target_selected_option]
        preprocess = st.selectbox('data preprocess', [None, 'StandardScaler', 'MinMaxScaler'])
        if preprocess == 'StandardScaler':
            features = StandardScaler().fit_transform(features)
        elif preprocess == 'MinMaxScaler':
            features = MinMaxScaler().fit_transform(features)

        model = pickle.load(model_file)
        prediction = model.predict(features)
        # st.write(std)
        plot = customPlot()
        plot.pred_vs_actual(targets, prediction)
        r2 = r2_score(targets, prediction)
        st.write('R2: {}'.format(r2))
        result_data = pd.concat([targets, pd.DataFrame(prediction)], axis=1)
        result_data.columns = ['actual', 'prediction']
        with st.expander('prediction'):
            st.write(result_data)
            tmp_download_link = download_button(result_data, f'prediction.csv', button_text='download')
            st.markdown(tmp_download_link, unsafe_allow_html=True)
        st.write('---')